import re
import socket
from sys import stdout

import paramiko

from common.libs.encrypt import decryption
from init_task.settings import AES_SECRET


def ssh_login(host, username, password):
    ssh = paramiko.SSHClient()
    ssh.set_missing_host_key_policy(paramiko.AutoAddPolicy())
    try:
        res, password = decryption(AES_SECRET, password)
        ssh.connect(hostname=host, username=username, password=password,timeout=10)
    except (socket.timeout, paramiko.ssh_exception.NoValidConnectionsError,paramiko.ssh_exception.AuthenticationException) as e:
        # 捕获连接失败、超时、认证失败等异常
        print(f"连接到{host} 失败:{str(e)}")
        return False,str(e)
    return True, ssh


def clean_output(output: str) -> str:
    # 匹配 ANSI 转义序列
    ansi_escape = re.compile(r'\x1B\[[0-?]*[ -/]*[@-~]')
    return ansi_escape.sub('',output)


def run_command(ssh: paramiko.SSHClient, command: str):
    if ssh.get_transport() is None or not ssh.get_transport().is_active():
        raise Exception("SSH连接无效，无法执行命令")
    stdin, stdout, stderr = ssh.exec_command(command)
    stdout_clean = clean_output(stdout.read().decode())
    return stdout_clean, stderr.read().decode()


def is_process_finished(ssh: paramiko.SSHClient, process_name: str):
    stdout, _ = run_command(ssh, f"pgrep -f '{process_name}'")
    return not bool(stdout.strip())